mod fem;
mod point_render;

use std::sync::Arc;
use std::fs::File;
use std::io::Read;
use std::time::Instant;

use fem::Fem;
use nalgebra::{vector, Vector2};
use point_render::PointRender;
use wgpu_text::{BrushBuilder, TextBrush};
use wgpu_widgets::d2::HeightCamera;
use winit::{
    application::ApplicationHandler, dpi::{PhysicalSize, Size}, event::{ElementState, KeyEvent, MouseButton, WindowEvent}, event_loop::EventLoop, keyboard::{KeyCode, PhysicalKey}, window::Window
};

fn main() {
    let mut app = App::new();
    let event_loop = EventLoop::new().unwrap();
    event_loop.run_app(&mut app).unwrap();
}

struct App<'a> {
    window: Option<Arc<Window>>,
    state: Option<State<'a>>,
}
impl App<'_> {
    fn new() -> Self {
        Self {
            window: None,
            state: None,
        }
    }
}

impl ApplicationHandler for App<'_> {
    fn resumed(&mut self, event_loop: &winit::event_loop::ActiveEventLoop) {
        let window = event_loop.create_window(
            Window::default_attributes()
            .with_inner_size(Size::Physical(PhysicalSize::new(512, 512)))
        ).unwrap();
        let window = Arc::new(window);
        self.window = Some(window.clone());
        self.state = Some(pollster::block_on(State::new(window)));
    }
    fn window_event(
        &mut self,
        event_loop: &winit::event_loop::ActiveEventLoop,
        _window_id: winit::window::WindowId,
        event: winit::event::WindowEvent,
    ) {
        match event {
            WindowEvent::CloseRequested => {
                println!("The close button was pressed; stopping");
                event_loop.exit();
            },
            WindowEvent::RedrawRequested => {
                let window = self.window.as_ref().unwrap();
                window.request_redraw();
                let state = self.state.as_mut().unwrap();
                state.render();
            },
            WindowEvent::Resized(size) => {
                let state = self.state.as_mut().unwrap();
                state.resize(size.width, size.height);
            },
            WindowEvent::MouseInput { button, state, .. } => {
                let s = self.state.as_mut().unwrap();
                s.on_mouse(button, state);
            },
            WindowEvent::CursorMoved { position, .. } => {
                let s = self.state.as_mut().unwrap();
                s.cursor_move(position.x, position.y);
            },
            WindowEvent::KeyboardInput { event, .. } => {
                let s = self.state.as_mut().unwrap();
                s.on_key(event);
            }
            _ => (),
        }
    }
}

fn read_file_to_string(path: &str) -> String {
    let mut file = File::open(path).unwrap();
    let mut s = String::new();
    file.read_to_string(&mut s).unwrap();
    s
}

fn read_file_to_vec(path: &str) -> Vec<u8> {
    let mut file = File::open(path).unwrap();
    let mut v = Vec::new();
    file.read_to_end(&mut v).unwrap();
    v
}

fn create_shader_from_string(device: &wgpu::Device, label: Option<&str>, s: String) -> wgpu::ShaderModule {
    device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label,
        source: wgpu::ShaderSource::Wgsl(s.into())
    })
}

fn create_shader_from_file(device: &wgpu::Device, path: &str) -> wgpu::ShaderModule {
    let src = read_file_to_string(path);
    create_shader_from_string(device, Some(path), src)
}

fn byte_size<T>(a: &[T]) -> usize {
    std::mem::size_of::<T>() * a.len()
}

unsafe fn cast_slice<T>(a: &[T]) -> &[u8] {
    let size = byte_size(a);
    let ptr = a as *const _ as *const u8;
    std::slice::from_raw_parts(ptr, size)
}

struct State<'a> {
    surface: wgpu::Surface<'a>,
    adapter: wgpu::Adapter,
    device: wgpu::Device,
    queue: wgpu::Queue,

    window_w: u32,
    window_h: u32,
    cursor_x: f64,
    cursor_y: f64,
    paused: bool,

    text_brush: TextBrush,
    widgets_renderer: wgpu_widgets::Renderer2D,
    point_render: PointRender,
    fem: Fem,

    debug_accumulate: Vec<f32>,
    debug_values: Vec<f32>,
    last_time: Instant,
    start_time: Instant,
    debug_interval: f64,
    frame_count: u64,
    fps: f64,
}

impl State<'_> {
    async fn new(window: Arc<Window>) -> Self {
        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::default());
        let size = window.inner_size();
        let surface = instance.create_surface(window).unwrap();
        let adapter = instance.request_adapter(
            &wgpu::RequestAdapterOptions {
                compatible_surface: Some(&surface),
                ..Default::default()
            }
        ).await.unwrap();
        let (device, queue) = adapter.request_device(
            &wgpu::DeviceDescriptor::default(),
            None,
        ).await.unwrap();
        let config = surface.get_default_config(
            &adapter,
            size.width,
            size.height,
        ).unwrap();

        let font = wgpu_text::glyph_brush::ab_glyph::FontArc::try_from_vec(read_file_to_vec("fonts/LiberationMono-Regular.ttf")).unwrap();
        let text_brush = BrushBuilder::using_font(font)
            .build(&device, size.width, size.height, config.format);

        let mut widgets_renderer = wgpu_widgets::Renderer2D::new(
            &device,
            size.width as f32 / size.height as f32,
            config.format
        );
        widgets_renderer.camera.pos = [0.5, 0.5];
        widgets_renderer.camera.height = 1.1;

        let point_render = PointRender::new(
            &widgets_renderer.camera_buffer,
            &device,
            config.format,
            0.005,
        );
        let fem = Fem::new();

        let n_debug = 2;
        let debug_values = vec![0.0; n_debug];
        let debug_accumulate = vec![0.0; n_debug];

        let time = Instant::now();
        Self {
            device, adapter, surface, queue,
            window_w: size.width, window_h: size.height,
            cursor_x: 0.0, cursor_y: 0.0, paused: true,
            text_brush, widgets_renderer, point_render, fem,
            debug_accumulate, debug_values,
            last_time: time, start_time: time, frame_count: 0, debug_interval: 0.5, fps: 0.0,
        }
    }
    fn render(&mut self) {
        let now = Instant::now();
        let dt = now - self.last_time;
        self.debug_accumulate[0] += dt.as_secs_f32();

        let debug_count = ((now - self.start_time).as_secs_f64() / self.debug_interval) as u64;
        let last_debug_count = ((self.last_time - self.start_time).as_secs_f64() / self.debug_interval) as u64;
        if debug_count > last_debug_count {
            for i in 0..self.debug_accumulate.len() {
                self.debug_values[i] = self.debug_accumulate[i] / self.frame_count as f32;
                self.debug_accumulate[i] = 0.0;
            }

            self.fps = self.frame_count as f64 / self.debug_interval;
            self.frame_count = 0;
        }
        self.frame_count += 1;
        
        // fem step
        if !self.paused {
            let t = Instant::now();

            let cal_dt = 1e-3;
            let step_count = 5.0;
            let cal_count = ((now - self.start_time).as_secs_f32() / cal_dt) as u64;
            let last_cal_count = ((self.last_time - self.start_time).as_secs_f32() / cal_dt) as u64;
            if cal_count > last_cal_count {
                let step_dt = cal_dt / step_count;
                for _ in 0..step_count as i32 {
                    self.fem.step(step_dt);
                }
            }
            // self.fem.step(3e-4);
            self.debug_values[1] = t.elapsed().as_secs_f32();
        }

        // debug info
        const DEBUG_COLOR: [f32; 4] = [0.0, 1.0, 0.0, 1.0];
        let debug_text = format!(
"fps: {:.2}
dt: {:.5}
cal_dt: {:.5}
cursor_force: {}
cursor_pos: {:?}",
            self.fps, self.debug_values[0], self.debug_values[1],
            self.fem.cursor_force,
            self.fem.cursor_world_pos,
        );
        self.text_brush.queue(
            &self.device,
            &self.queue,
            vec![wgpu_text::glyph_brush::Section::new()
                .add_text(wgpu_text::glyph_brush::Text::new(&debug_text).with_color(DEBUG_COLOR))
            ]
        ).unwrap();

        let mut encoder = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("encoder"),
        });
        let surface_texture = self.surface.get_current_texture().unwrap();
        {
            let surface_texture_view = surface_texture.texture.create_view(&Default::default());
            let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("render_pass"),
                color_attachments: &[
                    Some(wgpu::RenderPassColorAttachment {
                        view: &surface_texture_view,
                        resolve_target: None,
                        ops: wgpu::Operations {
                            load: wgpu::LoadOp::Clear(
                                wgpu::Color::BLACK
                            ),
                            store: wgpu::StoreOp::Store,
                        }
                    })
                ],
                depth_stencil_attachment: None,
                timestamp_writes: None,
                occlusion_query_set: None
            });
            self.point_render.update(&self.device, &self.queue, &self.fem.pxs);
            self.point_render.draw(&mut render_pass);
            self.widgets_renderer.draw(&mut render_pass);
            self.text_brush.draw(&mut render_pass);
        }
        self.queue.submit(std::iter::once(encoder.finish()));
        surface_texture.present();
        self.last_time = now;
    }
    fn resize(&mut self, width: u32, height: u32) {
        self.window_h = height;
        self.window_w = width;
        let config = self.surface.get_default_config(
            &self.adapter,
            width,
            height,
        ).unwrap();
        self.surface.configure(&self.device, &config);
        self.widgets_renderer.resize(width as f32 / height as f32, &self.queue);
        self.text_brush.resize_view(width as f32, height as f32, &self.queue);
    }
    fn on_mouse(&mut self, button: MouseButton, state: ElementState) {
        if button == MouseButton::Left {
            if state == ElementState::Pressed {
            }
        } else if button == MouseButton::Right {
            if state == ElementState::Pressed {
                self.fem.cursor_force = 30.0;
            } else {
                self.fem.cursor_force = 0.0;
            }
        }
    }
    fn cursor_move(&mut self, x: f64, y: f64) {
        self.cursor_x = x;
        self.cursor_y = y;
        self.fem.cursor_world_pos = screen_to_world(
            x, y,
            self.widgets_renderer.camera,
            self.window_w,
            self.window_h,
        );
    }
    fn on_key(&mut self, event: KeyEvent) {
        let keycode = if let PhysicalKey::Code(code) = event.physical_key {
            code
        } else {
            return;
        };
        if event.state == ElementState::Pressed {
            if keycode == KeyCode::Space {
                self.paused = !self.paused;
            } else if keycode == KeyCode::ArrowRight {
                self.fem.step(0.001);
            } else if keycode == KeyCode::ArrowLeft {
                self.fem.step(-0.001);
            }
        }
    }
}
fn screen_to_world(
    x: f64, y: f64,
    camera: HeightCamera,
    width: u32, height: u32,
) -> Vector2<f32> {
    let camera_width = camera.height / height as f32 * width as f32;
    vector!(
        (x as f32 / width as f32 * 2.0 - 1.0) * camera_width * 0.5 + camera.pos[0],
        (1.0 - y as f32 / height as f32 * 2.0) * camera.height * 0.5 + camera.pos[1],
    )
}